#!/usr/bin/python

from twisted.protocols.basic import Int16StringReceiver
from twisted.internet.protocol import Factory
from twisted.internet.endpoints import TCP4ClientEndpoint
from twisted.internet import reactor, defer

import hashlib

import srp

from   srp_auth import ocrypt
from   srp_auth import srp_auth_pb2

class InvalidMessage (Exception):
    pass    

class AuthenticationFailed (Exception):
    pass


def assert_fields(msg, required_fields):
    for f in required_fields.split():
        if f:
            if not msg.HasField(f):
                raise InvalidMessage

def phex(s):
    return ''.join( ['%x' % ord(c) for c in s] )


class SrpAuthFetchPublicInfoProtocol (Int16StringReceiver):
    
    
    def connectionMade(self):

        self.d = defer.Deferred()
        
        msg          = srp_auth_pb2.SRPAuthMessage()
        msg.msg_type = srp_auth_pb2.GET_PUBLIC_INFO
        
        self.sendString( msg.SerializeToString() )

        self.factory.set_connected(self)


    def stringReceived(self, s):
        if self.d:
            reply = srp_auth_pb2.SRPAuthMessage()
            reply.ParseFromString( s )
            
            if not (reply.msg_type == srp_auth_pb2.REPLY and reply.HasField('public_info')):
                self.d.errback("Failed to retrieve public info")
            else:                
                self.d.callback( (reply.public_info.hash_alg, reply.public_info.prime_size, reply.public_info.public_key_der) )
            self.d = None
            self.transport.loseConnection()
        else:
            pass # Ignore everything else

            
    # Retrurns deferred to (hash_alg, prime_size, public_key_der)
    def getPublicInfo(self):
        return self.d
            

    
class SrpFetchPublicInfoFactory (Factory):
    protocol = SrpAuthFetchPublicInfoProtocol
    def __init__(self):
        self.d_connected = defer.Deferred()
        
    def set_connected(self, srp_client):
        if self.d_connected:
            self.d_connected.callback(srp_client)
            self.d_connected = None
        
        
def fetchSRPServerPublicInfo( host, port ):
    f = SrpFetchPublicInfoFactory()

    f.d_connected.addCallbacks( lambda srp_client : srp_client.getPublicInfo() )

    ep = TCP4ClientEndpoint(reactor, host, port)

    ep.connect( f ).addErrback( f.d_connected.errback )

    return f.d_connected


            
class SrpAuthClientProtocol (Int16StringReceiver):
    
    
    def connectionMade(self):

        msg          = srp_auth_pb2.SRPAuthMessage()
        msg.msg_type = srp_auth_pb2.AUTHENTICATION
        
        tosvr = msg.auth

        self.usr = srp.User( self.factory.username, self.factory.password,
                             hash_alg = self.factory.hash_type,
                             ng_type  = self.factory.prime_type )
        
        tosvr.username, tosvr.A = self.usr.start_authentication()
        
        self.sendString( msg.SerializeToString() )
        
        self.step            = 1
        self.authenticated   = False
        self.request_queue   = list()
        self.req_in_progress = False
        self.srp_key         = None
    
        
    def _next_request(self):
        if self.authenticated and self.request_queue and not self.req_in_progress:
            self.sendString( self.request_queue[0][0].SerializeToString() )
            self.req_in_progress = True
            
            
    def _queue_request(self, msg, reply_handler):
        d = defer.Deferred()
        self.request_queue.append( (msg,d, reply_handler) )
        self._next_request()
        return d
            
            
    def _finish_request(self, reply):
        msg, d, reply_handler = self.request_queue[0]
        del self.request_queue[0]
        self.req_in_progress = False
        self._next_request()
        
        reply_handler( reply, d )

        
        
    def stringReceived(self, s):
        try:
            msg = srp_auth_pb2.SRPAuthMessage()
            msg.ParseFromString( s )
            
            if not self.authenticated:
                if not (msg.msg_type == srp_auth_pb2.AUTHENTICATION and msg.HasField('auth')):
                    raise AuthenticationFailed
                self.handleAuthMessage( msg )
            else:
                self._finish_request(msg)
                
        except InvalidMessage:
            print 'Invalid Message!'
            self.factory.set_connected(self, Exception('Invalid Network Message'))
            self.transport.loseConnection()
            
        except AuthenticationFailed:
            self.factory.set_connected(self, Exception('Invalid username/password'))
            self.transport.loseConnection()
        
            
            
    def handleAuthMessage(self, msg):
        auth = msg.auth        
        
        if self.step == 1:

            assert_fields( auth, 's B' )
            
            tosvr = srp_auth_pb2.SRPAuthMessage()
            tosvr.msg_type = srp_auth_pb2.AUTHENTICATION
            
            M = self.usr.process_challenge( auth.s, auth.B )
            
            if M is None:
                raise AuthenticationFailed
            
            tosvr.auth.M = M
            
            self.sendString( tosvr.SerializeToString() )
            
            self.step = 2
            
        elif self.step == 2:
            
            assert_fields( auth, 'HAMK' )
            
            self.usr.verify_session( auth.HAMK )
            
            if self.usr.authenticated():
                self.authenticated = True
                self.srp_key       = self.usr.get_session_key()
                self.factory.set_connected(self)
                self._next_request()
            else:
                raise AuthenticationFailed
                

        
    def _std_handler(self, reply, d):
        if reply.msg_type != srp_auth_pb2.REPLY:
            d.errback( Exception('Invalid Protocol') )
        elif reply.HasField('error_message'):
            d.errback( Exception(reply.error_message) )
        else:
            d.callback( None )
                
                
    def add_user(self, username, password, is_locked = False, expires_at = 0, is_admin = False):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.ADD_USER
        
        au = tosvr.add_user
        
        up = au.properties
        
        salt, verification_key = srp.create_salted_verification_key( username, password,
                                                                     self.factory.hash_type,
                                                                     self.factory.prime_type )
        
        au.salt             = salt
        au.verification_key = ocrypt.aes_oneshot_encrypt( self.srp_key, salt, verification_key )
        
        up.username   = username
        up.is_locked  = is_locked
        up.expires_at = expires_at
        up.is_admin   = is_admin
        
        return self._queue_request( tosvr, self._std_handler )
        
        
    def remove_user(self, username):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.REMOVE_USER
        
        tosvr.username = username
        
        return self._queue_request( tosvr, self._std_handler )
    
    
    def change_password(self, username, new_password):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.CHANGE_PASSWORD

        salt, verification_key = srp.create_salted_verification_key( username, new_password,
                                                                     self.factory.hash_type,
                                                                     self.factory.prime_type )
        
        cp = tosvr.change_password
        
        cp.username        = username
        cp.salt            = salt

        cp.verification_key = ocrypt.aes_oneshot_encrypt( self.srp_key, salt, verification_key )
        
        return self._queue_request( tosvr, self._std_handler )
        
        
    def get_users(self):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.GET_USER_LIST
        
        def reply_handler( reply, d ):
            if reply.msg_type != srp_auth_pb2.REPLY:
                d.errback( Exception('Invalid Protocol') )
            elif reply.HasField('error_message'):
                d.errback( Exception(reply.error_message) )
            elif not reply.user_list:
                d.errback( Exception('Missing user list') )
            else:
                d.callback( reply.user_list )
        
        return self._queue_request( tosvr, reply_handler )
        
        
    def get_user_properties(self, username):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.GET_USER_PROPERTIES
        
        tosvr.username = username
        
        def reply_handler( reply, d ):
            if reply.msg_type != srp_auth_pb2.REPLY:
                d.errback( Exception('Invalid Protocol') )
            elif reply.HasField('error_message'):
                d.errback( Exception(reply.error_message) )
            elif not reply.HasField('user_properties'):
                d.errback( Exception('Missing user properties') )
            else:
                d.callback( reply.user_properties )
        
        return self._queue_request( tosvr, reply_handler )
        
        
    def set_user_lock(self, username, is_locked):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.SET_USER_PROPERTIES
        
        tosvr.user_properties.username  = username
        tosvr.user_properties.is_locked = is_locked
        
        return self._queue_request( tosvr, self._std_handler )
        
        
    def set_user_admin(self, username, is_admin):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.SET_USER_PROPERTIES
        
        tosvr.user_properties.username = username
        tosvr.user_properties.is_admin = is_admin
        
        return self._queue_request( tosvr, self._std_handler )
        
        
    def set_user_expiration(self, username, expires_at):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.SET_USER_PROPERTIES
        
        tosvr.user_properties.username   = username
        tosvr.user_properties.expires_at = expires_at
        
        return self._queue_request( tosvr, self._std_handler )


    # returns deferred to (hash_alg, prime_size, public_key_der)
    # where hash alg = 0-5, and prime_size = 0-4 (corresponds to pysrp enumerations)
    def get_public_info(self):
        tosvr          = srp_auth_pb2.SRPAuthMessage()
        tosvr.msg_type = srp_auth_pb2.GET_PUBLIC_INFO
        
        def reply_handler( reply, d ):
            if reply.msg_type != srp_auth_pb2.REPLY:
                d.errback( Exception('Invalid Protocol') )
            elif reply.HasField('error_message'):
                d.errback( Exception(reply.error_message) )
            elif not reply.public_info.public_key_der:
                d.errback( Exception('Missing public key') )
            else:
                d.callback( (reply.public_info.hash_alg, reply.public_info.prime_size, reply.public_info.public_key_der) )
        
        return self._queue_request( tosvr, reply_handler )
        
        
    
class SrpClientFactory (Factory):
    protocol = SrpAuthClientProtocol
    def __init__(self, username, password, hash_type, prime_type):
        self.username    = username
        self.password    = password
        self.hash_type   = hash_type
        self.prime_type  = prime_type
        self.d_connected = defer.Deferred()
        
    def set_connected(self, srp_client, err=None):
        if self.d_connected:
            if err is None:
                self.d_connected.callback(srp_client)
            else:
                self.d_connected.errback( err )
            self.d_connected = None
